/*============================ Locomotion =============================*/
/**
 * FSM State for robot locomotion. Manages the contact specific logic
 * and handles calling the interfaces to the controllers. This state
 * should be independent of controller, gait, and desired trajectory.
 */

#include "FSM_State_Locomotion.h"
#include <AdvanceControllers/WBC_Ctrl/LocomotionCtrl/LocomotionCtrl.hpp>
using namespace std;

void RPYLimit(Vec3<float>& rpy, float max, float threshold, float alpha = 1);
/**
 * Constructor for the FSM State that passes in state specific info to
 * the generic FSM State constructor.
 *
 * @param _controlFSMData holds all of the relevant control data
 */
template <typename T>
FSM_State_Locomotion<T>::FSM_State_Locomotion(ControlFSMData<T>* _controlFSMData)
    : FSM_State<T>(_controlFSMData, FSM_StateName::LOCOMOTION, "LOCOMOTION") {
    // setup mpc controller
    cMPCOld = new ConvexMPCLocomotion(
        _controlFSMData->controlParameters->controller_dt,
        80 / (1000. * _controlFSMData->controlParameters
                          ->controller_dt),    // 27/(1000. * _controlFSMData->controlParameters->controller_dt)
        _controlFSMData->userParameters);

    this->turnOnAllSafetyChecks();
    //  Turn off Foot pos command since it is set in WBC as operational task
    this->checkPDesFoot = false;

    // Initialize GRF and footstep locations to 0s
    this->footFeedForwardForces = Mat34<T>::Zero();
    this->footstepLocations = Mat34<T>::Zero();

    // wbc
    _wbc_ctrl = new LocomotionCtrl<T>(_controlFSMData->_quadruped->buildModel());
    _wbc_data = new LocomotionCtrlData<T>();

    // balance control
    _bc_ctrl = new BalanceController<T>();
    _bc_data = new BalanceCtrlData<T>();

    // sports control
    _sp_ctrl = new SportsController<T>(_controlFSMData);

}

template <typename T>
void FSM_State_Locomotion<T>::onEnter() {
    // Default is to not transition
    this->nextStateName = this->stateName;

    // Reset the transition data
    this->transitionData.zero();
    cMPCOld->initialize();
    printf("[FSM LOCOMOTION] On Enter\n");
}

/**
 * Calls the functions to be executed on each control loop iteration.
 */
template <typename T>
void FSM_State_Locomotion<T>::run() {
    // Call the locomotion control logic for this iteration
    LocomotionControlStep(Mode_Pos);
}

extern rc_control_settings rc_control;

/**
 * Manages which states can be transitioned into either by the user
 * commands or state event triggers.
 *
 * @return the enumerated FSM state name to transition into
 */
template <typename T>
FSM_StateName FSM_State_Locomotion<T>::checkTransition() {
    // Get the next state
    iter++;

    // Switch FSM control mode
    // locomotionSafe()
    if (true) {
        switch ((int)this->_data->controlParameters->control_mode) {
            case K_LOCOMOTION:
                break;

            case K_BALANCE_STAND:
                // Requested change to BALANCE_STAND
                this->nextStateName = FSM_StateName::BALANCE_STAND;

                // Transition time is immediate
                this->transitionDuration = 0.0;

                break;

            case K_PASSIVE:
                // Requested change to BALANCE_STAND
                this->nextStateName = FSM_StateName::PASSIVE;

                // Transition time is immediate
                this->transitionDuration = 0.0;

                break;

            case K_STAND_UP:
                this->nextStateName = FSM_StateName::STAND_UP;
                this->transitionDuration = 0.;
                break;

            case K_RECOVERY_STAND:
                this->nextStateName = FSM_StateName::RECOVERY_STAND;
                this->transitionDuration = 0.;
                break;

            case K_VISION:
                this->nextStateName = FSM_StateName::VISION;
                this->transitionDuration = 0.;
                break;

            default:
                std::cout << "[CONTROL FSM] Bad Request: Cannot transition from " << K_LOCOMOTION << " to "
                          << this->_data->controlParameters->control_mode << std::endl;
        }
    } else {
        // this->nextStateName = FSM_StateName::STAND_UP;
        // this->transitionDuration = 0.;
        // rc_control.mode = RC_mode::STAND_UP;
    }

    // Return the next state name to the FSM
    return this->nextStateName;
}

/**
 * Handles the actual transition for the robot between states.
 * Returns true when the transition is completed.
 *
 * @return true if transition is complete
 */
template <typename T>
TransitionData<T> FSM_State_Locomotion<T>::transition() {
    // Switch FSM control mode

    switch (this->nextStateName) {
        case FSM_StateName::BALANCE_STAND:
            LocomotionControlStep(Mode_Pos);

            iter++;
            if (iter >= this->transitionDuration * 1000) {
                this->transitionData.done = true;
            } else {
                this->transitionData.done = false;
            }

            break;

        case FSM_StateName::PASSIVE:
            this->turnOffAllSafetyChecks();

            this->transitionData.done = true;

            break;

        case FSM_StateName::STAND_UP:
            this->transitionData.done = true;
            break;

        case FSM_StateName::RECOVERY_STAND:
            this->transitionData.done = true;
            break;

        case FSM_StateName::VISION:
            this->transitionData.done = true;
            break;

        default:
            std::cout << "[CONTROL FSM] Something went wrong in transition" << std::endl;
    }

    // Return the transition data to the FSM
    return this->transitionData;
}

template <typename T>
bool FSM_State_Locomotion<T>::locomotionSafe() {
    auto seResult = this->_data->_stateEstimator->getResult();
    const T max_roll = 15;
    const T max_pitch = 15;

    if (std::fabs(seResult.rpy[0]) > ori::deg2rad(max_roll)) {
        printf("Unsafe locomotion: roll is %.3f degrees (max %.3f)\n", ori::rad2deg(seResult.rpy[0]), max_roll);
        return false;
    }

    if (std::fabs(seResult.rpy[1]) > ori::deg2rad(max_pitch)) {
        printf("Unsafe locomotion: pitch is %.3f degrees (max %.3f)\n", ori::rad2deg(seResult.rpy[1]), max_pitch);
        return false;
    }

    for (int leg = 0; leg < 4; leg++) {
        auto p_leg = this->_data->_legController->datas[leg].p;
        if (p_leg[2] > 0) {
            printf("Unsafe locomotion: leg %d is above hip (%.3f m)\n", leg, p_leg[2]);
            return false;
        }

        if (std::fabs(p_leg[1] > 0.18)) {
            printf("Unsafe locomotion: leg %d's y-position is bad (%.3f m)\n", leg, p_leg[1]);
            return false;
        }

        auto v_leg = this->_data->_legController->datas[leg].v.norm();
        if (std::fabs(v_leg) > 9.) {
            printf("Unsafe locomotion: leg %d is moving too quickly (%.3f m/s)\n", leg, v_leg);
            return false;
        }
    }
    return true;
}

/**
 * Cleans up the state information on exiting the state.
 */
template <typename T>
void FSM_State_Locomotion<T>::onExit() {
    // Nothing to clean up when exiting
    iter = 0;
}

/**
 * Calculate the commands for the leg controllers for each of the feet by
 * calling the appropriate balance controller and parsing the results for
 * each stance or swing leg.
 * @param mode 0 mixed controll     1 position control
 */
template <typename T>
void FSM_State_Locomotion<T>::LocomotionControlStep(Leg_Mode mode) {
    Vec3<T> pDes_backup[4];
    Vec3<T> vDes_backup[4];
    Mat3<T> Kp_backup[4];
    Mat3<T> Kd_backup[4];
    auto seResult = this->_data->_stateEstimator->getResult();
    /* location + force mixed controller */
    switch (mode) {
        case Mode_Tau:
            cMPCOld->run<T>(*this->_data);
            for (int leg(0); leg < 4; ++leg) {
                pDes_backup[leg] = this->_data->_legController->commands[leg].pDes;
                vDes_backup[leg] = this->_data->_legController->commands[leg].vDes;
                Kp_backup[leg] = this->_data->_legController->commands[leg].kpCartesian;
                Kd_backup[leg] = this->_data->_legController->commands[leg].kdCartesian;
            }

            if (this->_data->userParameters->use_wbc > 0.9) {
                _wbc_data->pBody_des = cMPCOld->pBody_des;
                _wbc_data->vBody_des = cMPCOld->vBody_des;
                _wbc_data->aBody_des = cMPCOld->aBody_des;

                _wbc_data->pBody_RPY_des = cMPCOld->pBody_RPY_des;
                _wbc_data->vBody_Ori_des = cMPCOld->vBody_Ori_des;

                for (size_t i(0); i < 4; ++i) {
                    _wbc_data->pFoot_des[i] = cMPCOld->pFoot_des[i];
                    _wbc_data->vFoot_des[i] = cMPCOld->vFoot_des[i];
                    _wbc_data->aFoot_des[i] = cMPCOld->aFoot_des[i];
                    _wbc_data->Fr_des[i] = cMPCOld->Fr_des[i];
                }
                _wbc_data->contact_state = cMPCOld->contact_state;
                _wbc_ctrl->run(_wbc_data, *this->_data);
            } else if (this->_data->userParameters->use_bc > 0.9) {
                _bc_data->pBody_des = cMPCOld->pBody_des;
                _bc_data->vBody_des = cMPCOld->vBody_des;
                _bc_data->aBody_des = cMPCOld->aBody_des;

                _bc_data->pBody_RPY_des = cMPCOld->pBody_RPY_des;
                _bc_data->vBody_Ori_des = cMPCOld->vBody_Ori_des;

                for (size_t i(0); i < 4; ++i) {
                    _bc_data->pFoot_des[i] = cMPCOld->pFoot_des[i];
                    _bc_data->vFoot_des[i] = cMPCOld->vFoot_des[i];
                    _bc_data->aFoot_des[i] = cMPCOld->aFoot_des[i];
                    _bc_data->Fr_des[i] = cMPCOld->Fr_des[i];
                }
                _bc_data->contact_state = cMPCOld->contact_state;
                _bc_ctrl->run(_bc_data, *this->_data);
            }

            for (int leg(0); leg < 4; ++leg) {
                // this->_data->_legController->commands[leg].pDes = pDes_backup[leg];
                this->_data->_legController->commands[leg].vDes = vDes_backup[leg];
                // this->_data->_legController->commands[leg].kpCartesian = Kp_backup[leg];
                this->_data->_legController->commands[leg].kdCartesian = Kd_backup[leg];
            }
            this->_data->_legController->setControlMode(All_Tau);
            this->_data->_legController->setEnabled(true);
            break;

        case Mode_Pos:
            _sp_ctrl->run();
            this->_data->_legController->setControlMode(All_Tau);
            // this->_data->_legController->setControlMode(All_Pos);
            this->_data->_legController->setEnabled(true);
            break;

        default:
            break;
    }
}

/**
 * Stance leg logic for impedance control. Prevent leg slipping and
 * bouncing, as well as tracking the foot velocity during high speeds.
 */
template <typename T>
void FSM_State_Locomotion<T>::StanceLegImpedanceControl(int leg) {
    // Impedance control for the stance leg
    this->cartesianImpedanceControl(leg, this->footstepLocations.col(leg), Vec3<T>::Zero(),
                                    this->_data->controlParameters->stand_kp_cartesian,
                                    this->_data->controlParameters->stand_kd_cartesian);
}

// template class FSM_State_Locomotion<double>;
template class FSM_State_Locomotion<float>;